#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
背景色滤除对话框模块
提供实时预览的背景色滤除功能
"""

import cv2
import numpy as np
from PyQt5.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QPushButton, 
                           QLabel, QSlider, QColorDialog, QGroupBox, QSpinBox, QCheckBox, QGridLayout, QLineEdit)
from PyQt5.QtCore import Qt, pyqtSignal
from PyQt5.QtGui import QImage, QPixmap, QColor, QPainter

class BackgroundRemovalDialog(QDialog):
    """背景滤除设置对话框"""
    
    def __init__(self, image, parent=None, bg_removal=None):
        """初始化对话框
        
        参数:
            image: 原始图像
            parent: 父窗口对象
            bg_removal: 背景滤除设置元组 (背景颜色列表, 待滤除颜色列表, 阈值)
        """
        super().__init__(parent)
        self.image = image
        self.parent = parent
        
        # 设置窗口属性
        self.setWindowTitle("背景滤除设置")
        self.setMinimumWidth(400)
        
        # 初始化颜色设置
        if bg_removal and len(bg_removal) >= 3:
            self.bg_color = bg_removal[0]  # 背景颜色
            self.target_color = bg_removal[1]  # 待滤除颜色
            # 确保阈值是整数
            self.threshold = int(bg_removal[2]) if isinstance(bg_removal[2], (int, float)) else 0
            # 如果有传入参数，则根据阈值决定是否启用滤除(阈值>0表示启用)
            self.enable_removal = self.threshold > 0
        else:
            # 背景颜色默认为绿色
            self.bg_color = [0, 255, 0]  # 绿色RGB值
            
            # 获取左上角第一个像素点的颜色作为默认待滤除颜色
            if image is not None and image.size > 0:
                self.target_color = [int(c) for c in image[0, 0]]
            else:
                self.target_color = [255, 255, 255]  # 白色
                
            self.threshold = 0  # 默认阈值设为0，表示不启用
            self.enable_removal = False  # 默认不启用滤除
        
        # 设置布局
        self.setup_ui()
        
        # 更新预览
        self.update_preview()
    
    def setup_ui(self):
        """设置UI界面"""
        layout = QVBoxLayout()
        
        # 背景颜色设置组
        bg_group = QGroupBox("背景颜色设置")
        bg_layout = QGridLayout()
        
        # 背景颜色选择按钮
        self.bg_color_button = QPushButton()
        self.bg_color_button.setFixedSize(30, 30)
        self.bg_color_button.clicked.connect(self.choose_bg_color)
        bg_layout.addWidget(self.bg_color_button, 0, 0)
        
        # 背景颜色RGB输入
        bg_layout.addWidget(QLabel("RGB:"), 0, 1)
        self.bg_r_spinbox = QSpinBox()
        self.bg_r_spinbox.setRange(0, 255)
        self.bg_r_spinbox.setValue(self.bg_color[0])
        self.bg_r_spinbox.valueChanged.connect(self.on_bg_rgb_changed)
        bg_layout.addWidget(self.bg_r_spinbox, 0, 2)
        
        self.bg_g_spinbox = QSpinBox()
        self.bg_g_spinbox.setRange(0, 255)
        self.bg_g_spinbox.setValue(self.bg_color[1])
        self.bg_g_spinbox.valueChanged.connect(self.on_bg_rgb_changed)
        bg_layout.addWidget(self.bg_g_spinbox, 0, 3)
        
        self.bg_b_spinbox = QSpinBox()
        self.bg_b_spinbox.setRange(0, 255)
        self.bg_b_spinbox.setValue(self.bg_color[2])
        self.bg_b_spinbox.valueChanged.connect(self.on_bg_rgb_changed)
        bg_layout.addWidget(self.bg_b_spinbox, 0, 4)
        
        bg_group.setLayout(bg_layout)
        layout.addWidget(bg_group)
        
        # 待滤除颜色设置组
        target_group = QGroupBox("待滤除颜色设置")
        target_layout = QGridLayout()
        
        # 待滤除颜色选择按钮
        self.target_color_button = QPushButton()
        self.target_color_button.setFixedSize(30, 30)
        self.target_color_button.clicked.connect(self.choose_target_color)
        target_layout.addWidget(self.target_color_button, 0, 0)
        
        # 待滤除颜色RGB输入
        target_layout.addWidget(QLabel("RGB:"), 0, 1)
        self.target_r_spinbox = QSpinBox()
        self.target_r_spinbox.setRange(0, 255)
        self.target_r_spinbox.setValue(self.target_color[0])
        self.target_r_spinbox.valueChanged.connect(self.on_target_rgb_changed)
        target_layout.addWidget(self.target_r_spinbox, 0, 2)
        
        self.target_g_spinbox = QSpinBox()
        self.target_g_spinbox.setRange(0, 255)
        self.target_g_spinbox.setValue(self.target_color[1])
        self.target_g_spinbox.valueChanged.connect(self.on_target_rgb_changed)
        target_layout.addWidget(self.target_g_spinbox, 0, 3)
        
        self.target_b_spinbox = QSpinBox()
        self.target_b_spinbox.setRange(0, 255)
        self.target_b_spinbox.setValue(self.target_color[2])
        self.target_b_spinbox.valueChanged.connect(self.on_target_rgb_changed)
        target_layout.addWidget(self.target_b_spinbox, 0, 4)
        
        # 从图像选择颜色按钮
        self.pick_color_button = QPushButton("从图像选择")
        self.pick_color_button.clicked.connect(self.on_pick_color_clicked)
        target_layout.addWidget(self.pick_color_button, 1, 0, 1, 5)
        
        target_group.setLayout(target_layout)
        layout.addWidget(target_group)
        
        # 滤除功能开关
        self.enable_removal_checkbox = QCheckBox("启用背景滤除")
        self.enable_removal_checkbox.setChecked(self.enable_removal)
        self.enable_removal_checkbox.toggled.connect(self.on_enable_removal_toggled)
        layout.addWidget(self.enable_removal_checkbox)
        
        # 滤除强度设置
        strength_group = QGroupBox("滤除强度")
        strength_layout = QVBoxLayout()
        
        self.strength_slider = QSlider(Qt.Horizontal)
        self.strength_slider.setRange(0, 100)
        self.strength_slider.setValue(self.threshold)
        self.strength_slider.valueChanged.connect(self.on_strength_changed)
        strength_layout.addWidget(self.strength_slider)
        
        self.strength_label = QLabel(f"强度: {self.threshold}")
        strength_layout.addWidget(self.strength_label)
        
        strength_group.setLayout(strength_layout)
        layout.addWidget(strength_group)
        
        # 预览区域
        preview_group = QGroupBox("预览")
        preview_layout = QVBoxLayout()
        
        self.preview_label = QLabel()
        self.preview_label.setMinimumSize(300, 200)
        self.preview_label.setAlignment(Qt.AlignCenter)
        self.preview_label.mousePressEvent = self.on_preview_clicked
        preview_layout.addWidget(self.preview_label)
        
        preview_group.setLayout(preview_layout)
        layout.addWidget(preview_group)
        
        # 按钮
        button_layout = QHBoxLayout()
        button_layout.addStretch()
        
        ok_button = QPushButton("确定")
        ok_button.clicked.connect(self.accept)
        button_layout.addWidget(ok_button)
        
        cancel_button = QPushButton("取消")
        cancel_button.clicked.connect(self.reject)
        button_layout.addWidget(cancel_button)
        
        layout.addLayout(button_layout)
        
        self.setLayout(layout)
        
        # 更新颜色按钮显示
        self.update_color_buttons()
        
        # 设置初始状态
        self.on_enable_removal_toggled(self.enable_removal)
        
        # 颜色选择模式标志
        self.is_picking_color = False
    
    def update_color_buttons(self):
        """更新颜色按钮的显示"""
        # 更新背景颜色按钮
        self.bg_color_button.setStyleSheet(
            f"background-color: rgb({self.bg_color[0]}, {self.bg_color[1]}, {self.bg_color[2]}); border: 1px solid black;"
        )
        
        # 更新待滤除颜色按钮
        self.target_color_button.setStyleSheet(
            f"background-color: rgb({self.target_color[0]}, {self.target_color[1]}, {self.target_color[2]}); border: 1px solid black;"
        )
    
    def choose_bg_color(self):
        """选择背景颜色"""
        color = QColorDialog.getColor(
            QColor(*self.bg_color),
            self,
            "选择背景颜色"
        )
        
        if color.isValid():
            self.bg_color = [color.red(), color.green(), color.blue()]
            self.bg_r_spinbox.setValue(self.bg_color[0])
            self.bg_g_spinbox.setValue(self.bg_color[1])
            self.bg_b_spinbox.setValue(self.bg_color[2])
            self.update_color_buttons()
            self.update_preview()
    
    def choose_target_color(self):
        """选择待滤除颜色"""
        color = QColorDialog.getColor(
            QColor(*self.target_color),
            self,
            "选择待滤除颜色"
        )
        
        if color.isValid():
            self.target_color = [color.red(), color.green(), color.blue()]
            self.target_r_spinbox.setValue(self.target_color[0])
            self.target_g_spinbox.setValue(self.target_color[1])
            self.target_b_spinbox.setValue(self.target_color[2])
            self.update_color_buttons()
            self.update_preview()
    
    def on_bg_rgb_changed(self):
        """背景颜色RGB值改变事件"""
        self.bg_color = [
            self.bg_r_spinbox.value(),
            self.bg_g_spinbox.value(),
            self.bg_b_spinbox.value()
        ]
        self.update_color_buttons()
        self.update_preview()
    
    def on_target_rgb_changed(self):
        """待滤除颜色RGB值改变事件"""
        self.target_color = [
            self.target_r_spinbox.value(),
            self.target_g_spinbox.value(),
            self.target_b_spinbox.value()
        ]
        self.update_color_buttons()
        self.update_preview()
    
    def on_strength_changed(self):
        """滤除强度改变事件"""
        self.threshold = self.strength_slider.value()
        self.strength_label.setText(f"强度: {self.threshold}")
        self.update_preview()
    
    def on_enable_removal_toggled(self, checked):
        """滤除功能开关改变事件"""
        self.enable_removal = checked
        # 启用/禁用相关控件
        self.strength_slider.setEnabled(checked)
        self.strength_label.setEnabled(checked)
        self.update_preview()
    
    def on_pick_color_clicked(self):
        """从图像选择颜色按钮点击事件"""
        self.is_picking_color = True
        self.pick_color_button.setText("请在预览区域点击选择颜色")
        self.preview_label.setCursor(Qt.CrossCursor)
    
    def on_preview_clicked(self, event):
        """预览区域点击事件处理"""
        if not self.is_picking_color or self.image is None:
            return
            
        # 获取点击位置
        pixmap = self.preview_label.pixmap()
        if pixmap is None:
            return
            
        # 计算点击位置在原图中的坐标
        label_width = self.preview_label.width()
        label_height = self.preview_label.height()
        pixmap_width = pixmap.width()
        pixmap_height = pixmap.height()
        
        # 计算缩放比例
        scale_x = pixmap_width / self.image.shape[1]
        scale_y = pixmap_height / self.image.shape[0]
        scale = min(scale_x, scale_y)
        
        # 计算图像在标签中的偏移
        offset_x = (label_width - pixmap_width) / 2
        offset_y = (label_height - pixmap_height) / 2
        
        # 计算点击位置在原图中的坐标
        x = int((event.x() - offset_x) / scale)
        y = int((event.y() - offset_y) / scale)
        
        # 确保坐标在图像范围内
        if 0 <= x < self.image.shape[1] and 0 <= y < self.image.shape[0]:
            # 获取点击位置的颜色
            color = self.image[y, x].tolist()
            self.target_color = color
            
            # 更新界面
            self.target_r_spinbox.setValue(color[0])
            self.target_g_spinbox.setValue(color[1])
            self.target_b_spinbox.setValue(color[2])
            self.update_color_buttons()
            self.update_preview()
            
            # 重置拾色状态
            self.is_picking_color = False
            self.preview_label.setCursor(Qt.ArrowCursor)
            self.pick_color_button.setText("从图像选择")
    
    def update_preview(self):
        """更新预览图像"""
        if self.image is None:
            return
            
        # 创建预览图像副本
        preview_img = self.image.copy()
        
        # 如果启用了滤除功能且阈值大于0，应用滤除效果
        if self.enable_removal and self.threshold > 0:
            # 计算颜色差异
            diff = np.abs(preview_img - np.array(self.target_color))
            mask = np.sum(diff, axis=2) <= self.threshold
            
            # 创建RGBA图像
            rgba = cv2.cvtColor(preview_img, cv2.COLOR_BGR2BGRA)
            
            # 只修改匹配像素的alpha通道为0（完全透明）
            rgba[mask, 3] = 0
            
            # 为了在预览中显示，我们需要将透明部分用背景色填充
            # 注意：这只是为了预览效果，实际处理时会保持透明
            bg = np.full_like(preview_img, self.bg_color)
            
            # 创建一个临时的RGB图像用于显示
            temp = rgba.copy()
            for i in range(3):  # 对BGR三个通道
                temp[mask, i] = self.bg_color[i]
            
            result = cv2.cvtColor(temp, cv2.COLOR_BGRA2BGR)
        else:
            # 如果未启用滤除功能或阈值为0，显示原图
            result = preview_img
        
        # 转换图像格式
        height, width = result.shape[:2]
        bytes_per_line = 3 * width
        q_img = QImage(
            result.data,
            width,
            height,
            bytes_per_line,
            QImage.Format_BGR888
        )  # OpenCV使用BGR顺序，QImage.Format_BGR888表示BGR顺序
        
        # 缩放图像以适应预览区域
        pixmap = QPixmap.fromImage(q_img)
        scaled_pixmap = pixmap.scaled(
            self.preview_label.size(),
            Qt.KeepAspectRatio,
            Qt.SmoothTransformation
        )
        
        # 显示预览
        self.preview_label.setPixmap(scaled_pixmap)
    
    def get_settings(self):
        """获取设置
        
        返回:
            tuple: (背景颜色列表, 待滤除颜色列表, 阈值)
        """
        return (self.bg_color, self.target_color, self.threshold)
    
    def showEvent(self, event):
        """显示事件处理"""
        super().showEvent(event)
        
        # 通知主窗口对话框已显示
        if self.parent is not None:
            self.parent.dialog_active = True
            if hasattr(self.parent, 'canvas') and hasattr(self.parent.canvas, 'pause_suggestion'):
                self.parent.canvas.pause_suggestion = True
    
    def closeEvent(self, event):
        """关闭事件处理"""
        # 通知主窗口对话框已关闭
        if self.parent is not None:
            self.parent.dialog_active = False
            if hasattr(self.parent, 'canvas') and hasattr(self.parent.canvas, 'pause_suggestion'):
                self.parent.canvas.pause_suggestion = False
        
        super().closeEvent(event)

    def on_color_picked(self, color):
        """处理从图像中选择的颜色
        
        参数:
            color (list): 选中的RGB颜色值
        """
        if color and len(color) >= 3:
            self.target_color = [int(c) for c in color[:3]]
            self.target_r_spinbox.setValue(self.target_color[0])
            self.target_g_spinbox.setValue(self.target_color[1])
            self.target_b_spinbox.setValue(self.target_color[2])
            self.update_color_buttons()
            self.update_preview()
            self.is_picking_color = False
            self.preview_label.setCursor(Qt.ArrowCursor)
            self.pick_color_button.setText("从图像选择") 